package com.allwees.core.mq.rabbitmq;

import lombok.extern.slf4j.Slf4j;
import org.springframework.amqp.core.AmqpTemplate;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.core.MessageDeliveryMode;
import org.springframework.amqp.core.MessageProperties;
import org.springframework.amqp.rabbit.connection.RabbitUtils;
import org.springframework.amqp.rabbit.core.RabbitTemplate;
import org.springframework.amqp.rabbit.retry.MessageRecoverer;
import org.springframework.util.Assert;

import java.io.PrintWriter;
import java.io.StringWriter;
import java.util.Map;

/**
 * 死信消息（达到最大重试次数）处理器
 *
 * @author yangylsky
 * @since 2020-03-31
 *
 */
@Slf4j
public class DeadLetterMessageRecoverer implements MessageRecoverer {
	public static final String X_EXCEPTION_STACKTRACE = "x-exception-stacktrace";

	public static final String X_EXCEPTION_MESSAGE = "x-exception-message";

	public static final String X_ORIGINAL_EXCHANGE = "x-original-exchange";

	public static final String X_ORIGINAL_ROUTING_KEY = "x-original-routingKey";

	public static final int DEFAULT_FRAME_MAX_HEADROOM = 20_000;

	private int frameMaxHeadroom = DEFAULT_FRAME_MAX_HEADROOM;

	protected final AmqpTemplate deadTemplate;

	protected final String deadRoutingKey;

	protected final String deadExchangeName;

	private String deadRoutingKeyPrefix = "dead.";

	private volatile Integer maxStackTraceLength = -1;

	private MessageDeliveryMode deliveryMode = MessageDeliveryMode.PERSISTENT;

	public DeadLetterMessageRecoverer(AmqpTemplate deadTemplate) {
		this(deadTemplate, null, null);
	}

	public DeadLetterMessageRecoverer(AmqpTemplate deadTemplate, String deadExchange) {
		this(deadTemplate, deadExchange, null);
	}

	public DeadLetterMessageRecoverer(AmqpTemplate deadTemplate, String deadExchange, String deadRoutingKey) {
		Assert.notNull(deadTemplate, "'deadTemplate' cannot be null");
		this.deadTemplate = deadTemplate;
		this.deadExchangeName = deadExchange;
		this.deadRoutingKey = deadRoutingKey;
		if (!(this.deadTemplate instanceof RabbitTemplate)) {
			this.maxStackTraceLength = Integer.MAX_VALUE;
		}
	}

	public DeadLetterMessageRecoverer deadRoutingKeyPrefix(String deadRoutingKeyPrefix) {
		this.setDeadRoutingKeyPrefix(deadRoutingKeyPrefix);
		return this;
	}

	public void setDeadRoutingKeyPrefix(String deadRoutingKeyPrefix) {
		Assert.notNull(deadRoutingKeyPrefix, "'deadRoutingKeyPrefix' cannot be null");
		this.deadRoutingKeyPrefix = deadRoutingKeyPrefix;
	}

	protected String getDeadRoutingKeyPrefix() {
		return this.deadRoutingKeyPrefix;
	}

	public void setDeliveryMode(MessageDeliveryMode deliveryMode) {
		Assert.notNull(deliveryMode, "'deliveryMode' cannot be null");
		this.deliveryMode = deliveryMode;
	}

	protected MessageDeliveryMode getDeliveryMode() {
		return this.deliveryMode;
	}

	@Override
	public void recover(Message message, Throwable cause) {
		MessageProperties messageProperties = message.getMessageProperties();
		Map<String, Object> headers = messageProperties.getHeaders();
		String stackTraceAsString = processStackTrace(cause);
		headers.put(X_EXCEPTION_STACKTRACE, stackTraceAsString);
		headers.put(X_EXCEPTION_MESSAGE, cause.getCause() != null ? cause.getCause().getMessage() : cause.getMessage());
		headers.put(X_ORIGINAL_EXCHANGE, messageProperties.getReceivedExchange());
		headers.put(X_ORIGINAL_ROUTING_KEY, messageProperties.getReceivedRoutingKey());
		Map<? extends String, ? extends Object> additionalHeaders = additionalHeaders(message, cause);
		if(additionalHeaders != null) {
			headers.putAll(additionalHeaders);
		}

		if(messageProperties.getDeliveryMode() == null) {
			messageProperties.setDeliveryMode(this.deliveryMode);
		}

		String routingKey = this.deadRoutingKey != null ? this.deadRoutingKey : this.prefixedOriginalRoutingKey(message);
		this.deadTemplate.send(this.deadExchangeName, routingKey, message);
		if(log.isWarnEnabled()) {
			log.warn("Republishing failed message to exchange '" + this.deadExchangeName + "' with routing key " + routingKey);
		}
	}

	private String processStackTrace(Throwable cause) {
		String stackTraceAsString = getStackTraceAsString(cause);
		if (this.maxStackTraceLength < 0) {
			int maxStackTraceLen = RabbitUtils
					.getMaxFrame(((RabbitTemplate) this.deadTemplate).getConnectionFactory());
			if (maxStackTraceLen > 0) {
				maxStackTraceLen -= this.frameMaxHeadroom;
				this.maxStackTraceLength = maxStackTraceLen;
			}
		}
		if (this.maxStackTraceLength > 0 && stackTraceAsString.length() > this.maxStackTraceLength) {
			stackTraceAsString = stackTraceAsString.substring(0, this.maxStackTraceLength);
			log.warn("Stack trace in republished message header truncated due to frame_max limitations; "
					+ "consider increasing frame_max on the broker or reduce the stack trace depth", cause);
		}
		return stackTraceAsString;
	}

	protected Map<? extends String, ? extends Object> additionalHeaders(Message message, Throwable cause) {
		return null;
	}

	private String prefixedOriginalRoutingKey(Message message) {
		return this.deadRoutingKeyPrefix + message.getMessageProperties().getReceivedRoutingKey();
	}

	private String getStackTraceAsString(Throwable cause) {
		StringWriter stringWriter = new StringWriter();
		PrintWriter printWriter = new PrintWriter(stringWriter, true);
		cause.printStackTrace(printWriter);
		return stringWriter.getBuffer().toString();
	}
}
